package indigo.mybatis;

import java.util.Map;

import org.apache.ibatis.builder.SqlSourceBuilder;
import org.apache.ibatis.builder.StaticSqlSource;
import org.apache.ibatis.builder.annotation.ProviderSqlSource;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.scripting.defaults.RawSqlSource;
import org.apache.ibatis.scripting.xmltags.DynamicContext;
import org.apache.ibatis.scripting.xmltags.DynamicSqlSource;
import org.apache.ibatis.scripting.xmltags.SqlNode;
import org.apache.ibatis.session.Configuration;

import com.google.common.collect.Maps;

public abstract class SqlSourceWrapper implements SqlSource {
	private final Map<String, Object> extraParameters = Maps.newHashMap();

	private SqlSourceWrapper() {
	}

	public static SqlSourceWrapper wrap(SqlSource original, Map<String, Object> parameters) {
		if (original instanceof DynamicSqlSource) {
			return new DynamicSqlSourceWrapper(original, parameters);
		} else if (original instanceof StaticSqlSource
				|| original instanceof StaticSqlSource
				|| original instanceof RawSqlSource
				|| original instanceof ProviderSqlSource) {
			return new SimpleSqlSourceWrapper(original, parameters);
		} else if (original instanceof SqlSourceWrapper) {
			SqlSourceWrapper wrapper = (SqlSourceWrapper) original;
			wrapper.getExtraParameterMap().putAll(parameters);
			return wrapper;
		} else {
			throw new IllegalStateException("Unsupported SqlSource type: " + original.getClass());
		}
	}

	public Map<String, Object> getExtraParameterMap() {
		return extraParameters;
	}

	private static class SimpleSqlSourceWrapper extends SqlSourceWrapper {
		private final SqlSource original;

		public SimpleSqlSourceWrapper(SqlSource original, Map<String, Object> extraParameters) {
			this.original = original;
			getExtraParameterMap().putAll(extraParameters);
		}

		@Override
		public BoundSql getBoundSql(Object parameterObject) {
			BoundSql boundSql = original.getBoundSql(parameterObject);
			for (Map.Entry<String, Object> entry : getExtraParameterMap().entrySet()) {
				boundSql.setAdditionalParameter(entry.getKey(), entry.getValue());
			}
			return boundSql;
		}
	}

	private static class DynamicSqlSourceWrapper extends SqlSourceWrapper {
		private final Configuration configuration;
		private final SqlNode rootSqlNode;

		public DynamicSqlSourceWrapper(SqlSource original, Map<String, Object> extraParameters) {
			MetaObject originalMeta = SystemMetaObject.forObject(original);
			this.configuration = (Configuration) originalMeta.getValue("configuration");
			this.rootSqlNode = (SqlNode) originalMeta.getValue("rootSqlNode");
			getExtraParameterMap().putAll(extraParameters);
		}

		@Override
		public BoundSql getBoundSql(Object parameterObject) {
			DynamicContext context = new DynamicContext(configuration, parameterObject);
			Map<String, Object> bindings = context.getBindings();
			bindings.putAll(getExtraParameterMap());
			rootSqlNode.apply(context);

			SqlSourceBuilder builder = new SqlSourceBuilder(configuration);
			Class<?> parameterType = parameterObject == null ? Object.class : parameterObject.getClass();
			SqlSource sqlSource = builder.parse(context.getSql(), parameterType, bindings);
			BoundSql boundSql = sqlSource.getBoundSql(parameterObject);
			for (Map.Entry<String, Object> entry : bindings.entrySet()) {
				boundSql.setAdditionalParameter(entry.getKey(), entry.getValue());
			}
			return boundSql;
		}
	}
}
